#coding=utf-8
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np


class WordEmbedding(nn.Module):
    """Word Embedding

    The ntoken-th dim is used for padding_idx, which agrees *implicitly*
    with the definition in Dictionary.
    """
    def __init__(self, ntoken, emb_dim, dropout): #for ntoken, emb_dim, dropout in zip(len(dictionary.word2idx), 300, 0.0)
        super(WordEmbedding, self).__init__()
        self.emb = nn.Embedding(ntoken+1, emb_dim, padding_idx=ntoken)
        '''
        一个保存了固定字典和大小的简单查找表。这个模块常用来保存词嵌入和用下标检索它们。模块的输入是一个下标的列表，输出是对应的词嵌入。

            参数：
            num_embeddings (int) - 嵌入字典的大小
            embedding_dim (int) - 每个嵌入向量的大小
            padding_idx (int, optional) - 如果提供的话，输出遇到此下标时用零填充
            max_norm (float, optional) - 如果提供的话，会重新归一化词嵌入，使它们的范数小于提供的值
            norm_type (float, optional) - 对于max_norm选项计算p范数时的p
            scale_grad_by_freq (boolean, optional) - 如果提供的话，会根据字典中单词频率缩放梯度
            
            变量：
            weight (Tensor) -形状为(num_embeddings, embedding_dim)的模块中可学习的权值
            
            形状：
            输入： LongTensor (N, W), N = mini-batch, W = 每个mini-batch中提取的下标数
            输出： (N, W, embedding_dim)
        '''
        self.dropout = nn.Dropout(dropout)
        ''' 避免循环神经网络过拟合，采用丢弃法。
            随机将输入张量中部分元素设置为0。对于每次前向调用，被置0的元素都是随机的。
            参数：
            p- 将元素置0 的概率；
            in-place 默认为False 如果设置为True就会原地操作。
            形状：
            输入：任意。 输入可以为任意形状；
            输出：相同。 输出和输入形状相同。
        '''

        self.ntoken = ntoken
        self.emb_dim = emb_dim

    def init_embedding(self, np_file): #np_file = data/glove6b_init_300d.npy
        weight_init = torch.from_numpy(np.load(np_file))
        assert weight_init.shape == (self.ntoken, self.emb_dim)
        self.emb.weight.data[:self.ntoken] = weight_init

    def forward(self, x): #x [batch_size, seq_length]
        emb = self.emb(x) #emb [batch_size, seq_length, 300]
        emb = self.dropout(emb) #emb [batch_size, seq_length, 300]
        return emb


class QuestionEmbedding(nn.Module): #for in_dim, num_hid, nlayers, bidirect, dropout in zip(300, num_hid(1024), 1, False, 0.0)
    def __init__(self, in_dim, num_hid, nlayers, bidirect, dropout, rnn_type='GRU'):
        """Module for question embedding
        """
        super(QuestionEmbedding, self).__init__()
        assert rnn_type == 'LSTM' or rnn_type == 'GRU'
        rnn_cls = nn.LSTM if rnn_type == 'LSTM' else nn.GRU #用的门控循环神经网络GRU

        self.rnn = rnn_cls(
            in_dim, num_hid, nlayers,
            bidirectional=bidirect,
            dropout=dropout,
            batch_first=True)
        '''
        input_size: 单个LSTM神经元的输入维度。
        hidden_size: 单个LSTM神经元的隐含层输出维度。
        num_layers: LSTM的层数，这里指的是叠起来的层数，而不是展开的层数，展开是自适应的。
        bias: 计算过程中是否需要偏置。
        batch_first: batch是否位于第一个维度，这里有一个很多人容易混淆的点，后面再进行说明。
        dropout: 中每一层的输出的dropout概率，默认为0即不进行dropout。需要注意的一点是最后一层的输出是不会加上dropout概率的。也就是说，当你只用到一层LSTM的时候，这个参数是不起作用的。在0.4.0中如果你只有一层并且dropout概率不为0的时候会抛出警告。
        bidirectional: 是否双向。当置为True的时候，输出会为将双向LSTM的输出进行拼接，输出的feature size会增加一倍。
        '''

        self.in_dim = in_dim
        self.num_hid = num_hid
        self.nlayers = nlayers
        self.rnn_type = rnn_type
        self.ndirections = 1 + int(bidirect)

    def init_hidden(self, batch):
        # just to get the type of tensor
        weight = next(self.parameters()).data
        hid_shape = (self.nlayers * self.ndirections, batch, self.num_hid)
        if self.rnn_type == 'LSTM':
            return (Variable(weight.new(*hid_shape).zero_()),
                    Variable(weight.new(*hid_shape).zero_()))
        else:
            return Variable(weight.new(*hid_shape).zero_())

    def forward(self, x):
        # x: [batch, sequence, in_dim(300)]
        batch = x.size(0)
        hidden = self.init_hidden(batch)
        self.rnn.flatten_parameters() #Resets parameter data pointer so that they can use faster code paths
        output, hidden = self.rnn(x, hidden)
        #output [sequence, batch, hidden_size * num_directions]
        #hidden [num_layers * num_directions, batch, hidden_size]
        if self.ndirections == 1:
            return output[:, -1] #[num_layers * num_directions, hidden_size]

        forward_ = output[:, -1, :self.num_hid]
        backward = output[:, 0, self.num_hid:]
        return torch.cat((forward_, backward), dim=1) #在给定维度上对输入的张量序列seq 进行连接操作

    def forward_all(self, x):
        # x: [batch, sequence, in_dim]
        batch = x.size(0)
        hidden = self.init_hidden(batch)
        self.rnn.flatten_parameters()
        output, hidden = self.rnn(x, hidden)
        return output

# if __name__ == '__main__':
#     we = WordEmbedding(10, 300, 0.0)
#     print(we)
#
#     qe = QuestionEmbedding(300, 1024, 1, False, 0.0)
#     print(qe)